#%%
import vtk
from vtk.util.numpy_support import vtk_to_numpy
import numpy as np
import os

def extract_pressure_from_vtk(vtk_file, output_npy_file, output_centroid_path):
    # Create a VTK reader for polydata
    reader = vtk.vtkPolyDataReader()
    reader.SetFileName(vtk_file)
    reader.Update()

    # Get the output of the reader
    polydata = reader.GetOutput()

    # Check for the existence of the 'p' array
    point_data = polydata.GetPointData()
    pressure_data = point_data.GetArray('p')
    if pressure_data is None:
        raise ValueError("Pressure data 'p' not found in the VTK file.")
    
    # Convert the pressure data to a numpy array
    pressure_values_np = vtk_to_numpy(pressure_data).astype(np.float64)
    
    # Get the number of cells
    num_cells = polydata.GetNumberOfCells()

    # Initialize a list to store cell pressure values
    cell_pressures = []
    cell_average_points = []

    for i in range(num_cells):
        # Get the cell
        cell = polydata.GetCell(i)
        
        # Get the points of the cell
        points = cell.GetPoints()
        
        # Calculate the average pressure for the cell
        cell_pressure = np.mean([pressure_values_np[cell.GetPointIds().GetId(j)] for j in range(points.GetNumberOfPoints())])
        cell_average_point = np.mean([points.GetPoint(j) for j in range(points.GetNumberOfPoints())], axis=0)
        
        # Append the cell pressure value
        cell_pressures.append(cell_pressure)
        cell_average_points.append(cell_average_point)
    
    # Convert the list of cell pressures to a numpy array
    cell_pressures_np = np.array(cell_pressures)
    cell_average_points_np = np.array(cell_average_points)

    # Save the cell pressure values as an npy file
    np.save(output_npy_file, cell_pressures_np)
    np.save(output_centroid_path, cell_average_points_np)


vtk_folder = 'vtk_dataset'
output_folder = 'npy_dataset'

# Ensure the output directory exists
os.makedirs(output_folder, exist_ok=True)

vtk_files = os.listdir(vtk_folder)
for file in vtk_files:
    vtk_path = os.path.join(vtk_folder, file)
    output_press_filename = 'press_'+file.lstrip('DrivAer_F_D_WM_WW_').rstrip('.vtk')+'.npy'
    output_centroid_filename = 'centroid_'+file.lstrip('DrivAer_F_D_WM_WW_').rstrip('.vtk')+'.npy'
    output_press_path = os.path.join(output_folder, output_press_filename)
    output_centroid_path = os.path.join(output_folder, output_centroid_filename)

    # Check if the VTK file exists before processing
    if os.path.isfile(vtk_path):
        try:
            extract_pressure_from_vtk(vtk_path, output_press_path, output_centroid_path)
            print(f'Successfully processed {vtk_path}')
        except ValueError as e:
            print(f'Error processing {vtk_path}: {e}')
    else:
        print(f'VTK file {vtk_path} does not exist.')


# my_press_folder = 'datasets/npyData'
# real_press_folder = 'datasets'
# for i in [2,3,4]:
#     my_press = np.load(os.path.join(my_press_folder, f'press_{i:04d}.npy'))
#     real_press = np.load(os.path.join(real_press_folder, f'press_{i:04d}.npy'))
#     my_centroid = np.load(os.path.join(my_press_folder, f'centroid_{i:04d}.npy'))
#     real_centroid = np.load(os.path.join(real_press_folder, f'centroid_{i:04d}.npy'))
#     error_p = np.mean(abs(my_press - real_press))
#     error_c = np.mean(abs(my_centroid - real_centroid))
#     if np.array_equal(my_press, real_press) and np.array_equal(my_centroid, real_centroid):
#         print("The same.")
#     else:
#         # raise ValueError("Different matrices found.")
#         print(f'Error of press_{i:04d} is: {error_p:.15f}')
#         print(f'Error of centroid_{i:04d} is: {error_c:.15f}')


# import numpy as np
# cen = np.load('datasets/npyData/centroid_0746.npy')
# press = np.load('datasets/npyData/press_0746.npy')
# print(cen.shape)
# print(press.shape)
